#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/FunctionInterfaces.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/MathExtras.h"
#include "mlir/Transforms/InliningUtils.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/AsmParser/Parser.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Type.h"
#include "llvm/Support/SourceMgr.h"

#include "hyper/Hyper.h"
#include <iostream>

using namespace mlir;
using namespace mlir::hyper;

#include "hyper/HyperDialect.cpp.inc"

//===----------------------------------------------------------------------===//
// Hyper dialect.
//===----------------------------------------------------------------------===//

void HyperDialect::initialize() {
  addTypes<AsyncTokenType>();
  addOperations<
#define GET_OP_LIST
#include "hyper/Hyper.cpp.inc"
      >();
  addAttributes<
#define GET_ATTRDEF_LIST
#include "hyper/HyperOpsAttributes.cpp.inc"
      >();
}

#include "hyper/HyperOpsEnums.cpp.inc"

#define GET_ATTRDEF_CLASSES
#include "hyper/HyperOpsAttributes.cpp.inc"

//===----------------------------------------------------------------------===//
// AsyncOpInterface
//===----------------------------------------------------------------------===//

void mlir::hyper::addAsyncDependency(Operation *op, Value token) {
  op->insertOperands(0, {token});
  if (!op->template hasTrait<OpTrait::AttrSizedOperandSegments>())
    return;
  auto attrName =
      OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr();
  auto sizeAttr = op->template getAttrOfType<DenseI32ArrayAttr>(attrName);

  // Async dependencies is the only variadic operand.
  if (!sizeAttr)
    return;

  SmallVector<int32_t, 8> sizes(sizeAttr.asArrayRef());
  ++sizes.front();
  op->setAttr(attrName, Builder(op->getContext()).getDenseI32ArrayAttr(sizes));
}

//===----------------------------------------------------------------------===//
// ForOp
//===----------------------------------------------------------------------===//

void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
                  Value ub, ValueRange memIns, ValueRange memOuts,
                  BodyBuilderFn bodyBuilder, ArrayAttr isSharedMemAttr) {
  result.addOperands({lb, ub});
  result.addOperands(memIns);
  result.addOperands(memOuts);
  Region *bodyRegion = result.addRegion();
  bodyRegion->push_back(new Block);
  Block &bodyBlock = bodyRegion->front();
  bodyBlock.addArgument(builder.getIndexType(), result.location);
  for (Value v : memIns)
    bodyBlock.addArgument(v.getType(), v.getLoc());
  for (Value v : memOuts)
    bodyBlock.addArgument(v.getType(), v.getLoc());

  // Create the default terminator if the builder is not provided and if the
  // iteration arguments are not provided. Otherwise, leave this to the caller
  // because we don't know which values to return from the loop.
  if (!bodyBuilder) {
    ForOp::ensureTerminator(*bodyRegion, builder, result.location);
  } else {
    OpBuilder::InsertionGuard guard(builder);
    builder.setInsertionPointToStart(&bodyBlock);
    bodyBuilder(
        builder, result.location, bodyBlock.getArgument(0),
        bodyBlock.getArguments().drop_front(1).drop_back(memOuts.size()),
        bodyBlock.getArguments().drop_front(1 + memIns.size()));
  }

  // Ensure valid `isSharedMem` attribute
  if (memIns.size() + memOuts.size() > 0) {
    if (!isSharedMemAttr) {
      // add default 'isSharedMem' attribute
      SmallVector<bool, 4> isShared;
      for (size_t i = 0; i < memIns.size() + memOuts.size(); i++) {
        isShared.push_back(false);
      }
      auto attr = builder.getBoolArrayAttr(isShared);
      result.addAttribute("isSharedMem", attr);
    } else {
      result.addAttribute("isSharedMem", isSharedMemAttr);
    }
  }

  // Ensure the 'operand_segment_sizes' attribute.
  result.addAttribute(ForOp::getOperandSegmentSizeAttr(),
                      builder.getDenseI32ArrayAttr(
                          {static_cast<int32_t>(1), static_cast<int32_t>(1),
                           static_cast<int32_t>(memIns.size()),
                           static_cast<int32_t>(memOuts.size())}));
}

LogicalResult ForOp::verifyRegions() {
  // Check that the body defines as single block argument for the induction
  // variable.
  auto *body = getBody();
  if (!body->getArgument(0).getType().isIndex())
    return emitOpError(
        "expected body first argument to be an index argument for "
        "the induction variable");

  return success();
}

/// Prints the initialization list in the form of
///   <prefix>(%inner = %outer, %inner2 = %outer2, <...>)
/// where 'inner' values are assumed to be region arguments and 'outer' values
/// are regular SSA values.
static void printInitializationList(OpAsmPrinter &p,
                                    Block::BlockArgListType blocksArgs,
                                    ValueRange initializers,
                                    StringRef prefix = "") {
  assert(blocksArgs.size() == initializers.size() &&
         "expected same length of arguments and initializers");
  if (initializers.empty())
    return;

  p << prefix << '(';
  llvm::interleaveComma(llvm::zip(blocksArgs, initializers), p, [&](auto it) {
    p << std::get<0>(it) << " = " << std::get<1>(it) << " : "
      << std::get<1>(it).getType();
  });
  p << ")";
}

void ForOp::print(OpAsmPrinter &p) {
  p << " " << getInductionVar() << " = " << getLowerBound() << " to "
    << getUpperBound();

  printInitializationList(p, getRegionMemInArgs(), getMemInOperands(),
                          " mem_in");
  printInitializationList(p, getRegionMemOutArgs(), getMemOutOperands(),
                          " mem_out");

  p.printRegion(getRegion(),
                /*printEntryBlockArgs=*/false,
                /*printBlockTerminators=*/false);
  p.printOptionalAttrDict((*this)->getAttrs());
}

ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
  auto &builder = parser.getBuilder();
  Type indexType = builder.getIndexType();

  OpAsmParser::Argument inductionVariable;
  inductionVariable.type = indexType;
  OpAsmParser::UnresolvedOperand lb, ub, memIn, memOut;

  // Parse the induction variable followed by '='.
  if (parser.parseArgument(inductionVariable) || parser.parseEqual() ||
      // Parse loop bounds.
      parser.parseOperand(lb) ||
      parser.resolveOperand(lb, indexType, result.operands) ||
      parser.parseKeyword("to") || parser.parseOperand(ub) ||
      parser.resolveOperand(ub, indexType, result.operands))
    return failure();

  // Parse the optional initial iteration arguments.
  SmallVector<OpAsmParser::Argument, 4> regionArgs;
  SmallVector<OpAsmParser::UnresolvedOperand, 4> operands;
  SmallVector<Type, 4> operandTypes;
  regionArgs.push_back(inductionVariable);

  // number of mem operands
  int numMemInOperands = 0;
  int numMemOutOperands = 0;

  // parse mem_in
  if (succeeded(parser.parseOptionalKeyword("mem_in"))) {
    // Parse assignment list and results type list.
    if (parser.parseLParen() ||
        parser.parseCommaSeparatedList([&]() -> ParseResult {
          if (parser.parseArgument(regionArgs.emplace_back()) ||
              parser.parseEqual() ||
              parser.parseOperand(operands.emplace_back()) ||
              parser.parseColonType(operandTypes.emplace_back()))
            return failure();
          return success();
        }) ||
        parser.parseRParen())
      return failure();

    // Resolve mem_in operands.
    int idx = 0;
    for (auto argOperandType :
         llvm::zip(llvm::drop_begin(regionArgs), operands, operandTypes)) {
      Type type = std::get<2>(argOperandType);
      auto operand = std::get<1>(argOperandType);
      std::get<0>(argOperandType).type = type;

      // type must be memref
      if (!type.isa<MemRefType>()) {
        return parser.emitError(parser.getNameLoc(),
                                "expected memref type for mem_in operand");
      }

      // the same memref operand can be used only once for mem_in
      for (int i = idx + 1; i < operands.size(); i++) {
        if (operand.name.equals(operands[i].name)) {
          return parser.emitError(
              parser.getNameLoc(),
              "memref operand can be used only once for mem_in");
        }
      }

      if (parser.resolveOperand(operand, type, result.operands))
        return failure();

      idx++;
    }
  }
  // number of mem_in operands
  numMemInOperands = regionArgs.size() - 1;
  // clear operands and types vectors
  operands.clear();
  operandTypes.clear();

  // parse mem_out
  if (succeeded(parser.parseOptionalKeyword("mem_out"))) {
    // Parse assignment list and results type list.
    if (parser.parseLParen() ||
        parser.parseCommaSeparatedList([&]() -> ParseResult {
          if (parser.parseArgument(regionArgs.emplace_back()) ||
              parser.parseEqual() ||
              parser.parseOperand(operands.emplace_back()) ||
              parser.parseColonType(operandTypes.emplace_back()))
            return failure();
          return success();
        }) ||
        parser.parseRParen())
      return failure();

    // Resolve mem_out operands.
    int idx = 0;
    for (auto argOperandType :
         llvm::zip(llvm::drop_begin(regionArgs, numMemInOperands + 1), operands,
                   operandTypes)) {
      Type type = std::get<2>(argOperandType);
      auto operand = std::get<1>(argOperandType);
      std::get<0>(argOperandType).type = type;

      // type must be memref
      if (!type.isa<MemRefType>()) {
        return parser.emitError(parser.getNameLoc(),
                                "expected memref type for mem_out operand");
      }

      // the same memref operand can be used only once for mem_out
      for (int i = idx + 1; i < operands.size(); i++) {
        if (operand.name.equals(operands[i].name)) {
          return parser.emitError(
              parser.getNameLoc(),
              "memref operand can be used only once for mem_out");
        }
      }

      if (parser.resolveOperand(operand, type, result.operands))
        return failure();

      idx++;
    }
  }
  // number of mem_out operands
  numMemOutOperands = regionArgs.size() - numMemInOperands - 1;

  // Parse the body region.
  Region *body = result.addRegion();
  if (parser.parseRegion(*body, regionArgs))
    return failure();

  ForOp::ensureTerminator(*body, builder, result.location);

  // Parse the optional attribute list.
  if (parser.parseOptionalAttrDict(result.attributes))
    return failure();

  int numMemOperands = numMemInOperands + numMemOutOperands;
  if (numMemOperands > 0) { // mem operands are present
    if (!result.attributes.getNamed("isSharedMem")
             .has_value()) { // 'isSharedMem' attribute is not present
      // add default 'isSharedMem' attribute
      SmallVector<bool, 4> isShared;
      for (size_t i = 0; i < numMemOperands; i++) {
        isShared.push_back(false);
      }
      result.addAttribute("isSharedMem", builder.getBoolArrayAttr(isShared));
    } else { // 'isSharedMem' attribute is present
      auto isSharedMemAttr = dyn_cast_or_null<ArrayAttr>(
          result.attributes.getNamed("isSharedMem").value().getValue());
      if (!isSharedMemAttr) { // check if the 'isSharedMem' attribute is of type
                              // 'array'
        return parser.emitError(
            parser.getNameLoc(),
            "expected 'isSharedMem' attribute to be of type 'array'");
      }
      if (isSharedMemAttr.size() !=
          numMemOperands) { // check if the size of 'isSharedMem' attribute is
                            // correct
        return parser.emitError(parser.getNameLoc(),
                                "expected 'isSharedMem' attribute to have size "
                                "equal to the number of mem operands");
      }
    }
  }

  // add operand segment size attribute
  if (!result.attributes.getNamed(getOperandSegmentSizeAttr()).has_value()) {
    result.addAttribute(ForOp::getOperandSegmentSizeAttr(),
                        builder.getDenseI32ArrayAttr(
                            {static_cast<int32_t>(1), static_cast<int32_t>(1),
                             static_cast<int32_t>(numMemInOperands),
                             static_cast<int32_t>(numMemOutOperands)}));
  }

  return success();
}

LogicalResult ForOp::verify() {
  auto isSharedMemAttr = getIsSharedMem();
  int numMemOperands = getNumMemVars();

  // check if the size of 'isSharedMem' attribute is correct
  if (!isSharedMemAttr.has_value()) {
    if (numMemOperands > 0) {
      return emitOpError("expected 'isSharedMem' attribute to be present");
    }
  } else {
    auto isSharedMem = isSharedMemAttr.value();
    if (isSharedMem.size() != numMemOperands) {
      return emitOpError("expected 'isSharedMem' attribute to have size equal "
                         "to the number of mem operands");
    }
  }

  // check the correctness of the 'devices' attribute
  auto devicesAttr = getDevices();
  float dutySum = 0.0;
  for (int i = 0; i < devicesAttr.size(); i++) {
    auto deviceDict = dyn_cast_or_null<DictionaryAttr>(devicesAttr[i]);

    // check if the main keys of the dictionary is present
    auto targetIdAttr = deviceDict.get("targetId");
    auto targetConfigAttr = deviceDict.get("targetConfig");
    auto dutyRatioAttr = deviceDict.get("dutyRatio");
    if (!targetIdAttr || !targetConfigAttr || !dutyRatioAttr) {
      return emitOpError("expected 'targetId', 'targetConfig' and 'duty_ratio' "
                         "attributes to be present in 'devices' attribute");
    }

    // check the types of each attribute in the dictionary
    auto targetId = targetIdAttr.dyn_cast_or_null<StringAttr>();
    if (!targetId) {
      return emitOpError(
          "expected 'targetId' attribute to be of type 'string'");
    }
    auto targetConfig = targetConfigAttr.dyn_cast_or_null<StringAttr>();
    if (!targetConfig) {
      return emitOpError(
          "expected 'targetConfig' attribute to be of type 'string'");
    }
    auto dutyRatio = dutyRatioAttr.dyn_cast_or_null<FloatAttr>();
    if (!dutyRatio) {
      return emitOpError(
          "expected 'dutyRatio' attribute to be of type 'float'");
    }

    // check if the targetId contains supported device
    if (!targetId.getValue().contains("cpu") &&
        !targetId.getValue().contains("gpu") &&
        !targetId.getValue().contains(
            "cgra")) { // only 'cpu' and 'gpu' are supported now
      return emitOpError("expected 'targetId' attribute to contain supported "
                         "device (be one of 'cpu', 'gpu' or 'cgra' for now)");
    }

    dutySum += dutyRatio.getValueAsDouble();
  }
  // sst don't need this condition
  // // check if the sum of duty ratios is equal to 1.0
  // if (dutySum != 1.0) {
  //   return emitOpError("expected the sum of duty ratios to be equal to 1.0");
  // }

  return success();
}

/// Given the region at `index`, or the parent operation if `index` is None,
/// return the successor regions. These are the regions that may be selected
/// during the flow of control. `operands` is a set of optional attributes that
/// correspond to a constant value for each operand, or null if that operand is
/// not a constant.
void ForOp::getSuccessorRegions(std::optional<unsigned> index,
                                ArrayRef<Attribute> operands,
                                SmallVectorImpl<RegionSuccessor> &regions) {
  // If the predecessor is the ForOp, branch into the body using the iterator
  // arguments.
  if (!index) {
    regions.push_back(RegionSuccessor(&getLoopBody()));
    return;
  }

  // Otherwise, the loop may branch back to itself or the parent operation.
  assert(*index == 0 && "expected loop region");
  regions.push_back(RegionSuccessor(&getLoopBody()));
}

Region &ForOp::getLoopBody() { return getRegion(); }

std::optional<Value> ForOp::getSingleInductionVar() {
  return getInductionVar();
}

std::optional<OpFoldResult> ForOp::getSingleLowerBound() {
  return OpFoldResult(getLowerBound());
}

std::optional<OpFoldResult> ForOp::getSingleUpperBound() {
  return OpFoldResult(getUpperBound());
}

//===----------------------------------------------------------------------===//
// ReduceOp
//===----------------------------------------------------------------------===//

void ReduceOp::build(
    OpBuilder &builder, OperationState &result, ValueRange operands,
    ArrayRef<function_ref<void(OpBuilder &, Location, Value, Value)>>
        bodyBuilderFns) {
  for (auto operandFn : llvm::zip(operands, bodyBuilderFns)) {
    auto operand = std::get<0>(operandFn);
    auto bodyBuilderFn = std::get<1>(operandFn);
    result.addOperands(operand);

    OpBuilder::InsertionGuard guard(builder);
    Region *bodyRegion = result.addRegion();
    auto type = operand.getType();
    Block *body =
        builder.createBlock(bodyRegion, {}, ArrayRef<Type>{type, type},
                            {result.location, result.location});

    if (bodyBuilderFn) {
      bodyBuilderFn(builder, result.location, body->getArgument(0),
                    body->getArgument(1));
    }
  }
}

LogicalResult ReduceOp::verifyRegions() {
  for (int i = 0; i < getReductionRegions().size(); i++) {
    // The region of a ReduceOp has two arguments of the same type as its
    // operand.
    auto type = getOperand(i).getType();
    Block &block = getReductionRegions()[i].front();
    if (block.empty())
      return emitOpError("the block inside reduce should not be empty");
    if (block.getNumArguments() != 2 ||
        llvm::any_of(block.getArguments(), [&](const BlockArgument &arg) {
          return arg.getType() != type;
        }))
      return emitOpError() << "expects two arguments to reduce block of type "
                           << type;

    // Check that the block is terminated by a ReduceReturnOp.
    if (!isa<ReduceReturnOp>(block.getTerminator()))
      return emitOpError("the block inside reduce should be terminated with a "
                         "'hyper.reduce.return' op");
  }

  return success();
}

LogicalResult ReduceOp::verify() {
  // The number of operands should be equal to the number of regions.
  if (getNumOperands() != getReductionRegions().size())
    return emitOpError(
        "expects the number of operands to be equal to the number of regions");

  // The type of operands should be as the same as shared mem_out.
  auto forOp = getParentOp();

  // Mem_out should be present if `hyper.reduce` op is used in a `hyper.for` op.
  if (!forOp.getIsSharedMem().has_value())
    return emitOpError("expected 'isSharedMem' attribute to be present in the "
                       "parent 'hyper.for' op");

  // Get the shared mem_outs array from the parent `hyper.for` op.
  SmallVector<Value> sharedMemOuts;
  for (auto memBool : llvm::zip(forOp.getMemOut(),
                                llvm::drop_begin(forOp.getIsSharedMem().value(),
                                                 forOp.getNumMemInVars()))) {
    auto mem = std::get<0>(memBool);
    auto isShared =
        std::get<1>(memBool).dyn_cast_or_null<BoolAttr>().getValue();
    if (isShared) {
      sharedMemOuts.push_back(mem);
    }
  }

  // The number of shared mem_outs should be equal to the number of operands.
  if (sharedMemOuts.size() != getNumOperands())
    return emitOpError("expected the number of shared mem_outs to be equal to "
                       "the number of operands");

  // The type of operands should be as the same as shared mem_out.
  for (auto operandMem : llvm::zip(getOperands(), sharedMemOuts)) {
    auto operand = std::get<0>(operandMem);
    auto mem = std::get<1>(operandMem);
    if (operand.getType() != mem.getType())
      return emitOpError("expected the type of operands to be equal to the "
                         "type of shared mem_outs");
  }

  return success();
}

ParseResult ReduceOp::parse(OpAsmParser &parser, OperationState &result) {
  // Parse an opening `(` followed by the comma-separated reduced values
  // followed by `)`
  if (parser.parseLParen() ||
      parser.parseCommaSeparatedList([&]() -> ParseResult {
        OpAsmParser::UnresolvedOperand operand;
        Type resultType;
        if (parser.parseOperand(operand) || parser.parseColonType(resultType) ||
            parser.resolveOperand(operand, resultType, result.operands))
          return failure();
        return success();
      }) ||
      parser.parseRParen())
    return failure();

  // Now parse comma-separated bodies.
  if (parser.parseCommaSeparatedList([&]() -> ParseResult {
        Region *body = result.addRegion();
        if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))
          return failure();
        return success();
      }))
    return failure();

  return success();
}

void ReduceOp::print(OpAsmPrinter &p) {
  p << "(";
  llvm::interleaveComma(getOperands(), p,
                        [&](auto it) { p << it << " : " << it.getType(); });
  p << ") ";

  for (unsigned i = 0; i < getReductionRegions().size(); i++) {
    p.printRegion(getReductionRegions()[i]);
    if (i < getReductionRegions().size() - 1) {
      p << ", ";
    }
  }
}

//===----------------------------------------------------------------------===//
// ReduceReturnOp
//===----------------------------------------------------------------------===//

LogicalResult ReduceReturnOp::verify() {
  // This op must return lhs operand of current block exactly.
  auto block = this->getOperation()->getBlock();
  auto lhs = block->getArguments()[0];
  auto retVal = getOperand();
  if (lhs != retVal)
    return emitOpError("must return lhs operand of current block exactly");

  return success();
}
